//  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
//  This source code is licensed under both the GPLv2 (found in the
//  COPYING file in the root directory) and Apache 2.0 License
//  (found in the LICENSE.Apache file in the root directory).
//
// Copyright (c) 2011 The LevelDB Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file. See the AUTHORS file for names of contributors.

#include "pure_mem/inlineart.h"
#include <set>
#include <iostream>
#include <unordered_set>
#include "rocksdb/env.h"
#include "util/concurrent_arena.h"
#include "util/hash.h"
#include "util/random.h"
#include "util/testharness.h"

namespace rocksdb {

// Our test skip list stores 8-byte unsigned integers
typedef uint64_t Key;

static const char* Encode(const Key* key) {
  std::string str = std::to_string(*key);
  int len = str.length() + 8 + 1;
  uint32_t keySize = VarintLength(len) + len;
  char* ret = new char[keySize];
  memset(ret , '\0', keySize);
  char*p = EncodeVarint32(ret, len);
  memcpy(p, str.c_str(), str.length());

  return reinterpret_cast<const char*>(ret);
}

static Key Decode(const char* key) {
  Slice buf = GetLengthPrefixedSlice(key);
  Slice mvccKey = Slice(buf.data(), buf.size() - 8);
  const char ts_size = mvccKey[mvccKey.size() - 1];
  if ((size_t)ts_size >= mvccKey.size()) {
    return 0;
  }
  long ret = std::atol(mvccKey.data());
  // std::cout << "decode for keyto : " << ret << ", mvcckey:" << mvccKey.ToString() <<std::endl;
  return ret;
}

struct TestComparator {
  typedef Key DecodedType;

  static DecodedType decode_key(const char* b) {
    return Decode(b);
  }

  int operator()(const char* a, const char* b) const {
    Slice sa = GetLengthPrefixedSlice(a);
    Slice sb = GetLengthPrefixedSlice(b);
    //std::cout << "sa: "<< sa.data() << ", size: " << sa.size() << ", sb: " << sb.data() << ", size: " << sb.size() <<std::endl;
    uint comparLength = std::min(sa.size(), sb.size());
    for (uint i =0; i < comparLength; i++){
      //std::cout << int(sa[i]) << "    " << int(sb[i]) << std::endl;
      if (sa[i] > sb[i])
        return 1;
      else if (sa[i] < sb[i])
        return -1;
    }
    if (sa.size() > sb.size())
      return 1;
    else if (sa.size() < sb.size())
      return -1;

    return 0;
  }

  int operator()(const char* a, const DecodedType b) const {
    if (Decode(a) < b) {
      return -1;
    } else if (Decode(a) > b) {
      return +1;
    } else {
      return 0;
    }
  }
};

typedef InlineART<TestComparator> TestInlineART;
char * genKey(InlineART<TestComparator> &list,Key key){
  std::string str = std::to_string(key);
  //std::cout << "gen key: " <<str<<std::endl;
  int len = str.length() + 8 + 1;
  uint32_t keySize = VarintLength(len) + len;
  char* ret = new char[keySize + VarintLength(len)];
  VersionNode *rett = new VersionNode();
  rett->CASSetKey(nullptr, ret);
  memset(ret , '\0', keySize);
  char*p = EncodeVarint32(ret, len);
  memcpy(p, str.c_str(), str.length());
  return (char*)rett;
}

class InlineARTTest : public testing::Test {
 public:
  void Insert(TestInlineART* list, Key key) {
    InsertWithHint(list, key, nullptr);
  }

  bool InsertWithHint(TestInlineART* list, Key key, void** hint) {
    keys_.insert(key);
    char* buf = genKey(*list, key);
    bool res = list->Insert(buf);
    return res;
  }

  void Validate(TestInlineART* list) {
    // Check keys exist.
    for (Key key : keys_) {
      ASSERT_TRUE(list->Contains(Encode(&key)));
    }
    // Iterate over the list, make sure keys appears in order and no extra
    // keys exist.
    TestInlineART::Iterator iter(list);
    ASSERT_FALSE(iter.Valid());
    Key zero = 0;
    iter.Seek(Encode(&zero));
    for (Key key : keys_) {
      ASSERT_TRUE(iter.Valid());
      ASSERT_EQ(key, Decode(iter.key()));
      iter.Next();
    }
    ASSERT_FALSE(iter.Valid());
  }

 private:
  std::set<Key> keys_;
};

TEST_F(InlineARTTest, Empty) {
  Arena arena;
  TestComparator cmp;
  InlineART<TestComparator> list(cmp);
  Key key = 10;
  ASSERT_TRUE(!list.Contains(Encode(&key)));
  InlineART<TestComparator>::Iterator iter(&list);
  ASSERT_TRUE(!iter.Valid());
  iter.SeekToFirst();
  ASSERT_TRUE(!iter.Valid());
  key = 100;
  iter.Seek(Encode(&key));
  ASSERT_TRUE(!iter.Valid());
  iter.SeekForPrev(Encode(&key));
  ASSERT_TRUE(!iter.Valid());
  iter.SeekToLast();
  ASSERT_TRUE(!iter.Valid());
}

TEST_F(InlineARTTest, InsertAndLookup) {
  const int N = 2000;
  const int R = 5000;
  Random rnd(1000);
  std::set<Key> keys;
  ConcurrentArena arena;
  TestComparator cmp;
  InlineART<TestComparator> list(cmp);

  for (int i = 0; i < N; i++) {
    Key key = rnd.Next() % R;
    std::cout << "index: " << i << ", key: " << key << std::endl;
    if (keys.insert(key).second) {
      char * ret = genKey(list, key);
      list.Insert(ret);
    }
  }
  std::cout <<" insert end. with num: " << keys.size() << std::endl;

  for (Key i = 0; i < R; i++) {
    //std::cout << "check contain. for key: " << i << std::endl;
    if (list.Contains(Encode(&i))) {
      ASSERT_EQ(keys.count(i), 1U);
    } else {
      ASSERT_EQ(keys.count(i), 0U);
    }
  }
  std::cout << "contain check success.." << std::endl;

  // iterator sum tests
  {
    InlineART<TestComparator>::Iterator iter(&list);
    ASSERT_TRUE(!iter.Valid());

    uint64_t zero = 0;
    iter.Seek(Encode(&zero));
    int listSum = 0;
    while (iter.Valid()){
      listSum++;
      //Slice buf = GetLengthPrefixedSlice(iter.key());
      //Slice mvccKey = Slice(buf.data(), buf.size() - 8);
      //std::cout <<"index : " << listSum << ", key: " << mvccKey.data() <<std::endl;
      iter.Next();
    }
    std::cout << "art leaf num: " << listSum << ", keys num: " << keys.size() <<std::endl;
    ASSERT_EQ(listSum, keys.size());
  }
  std::cout <<"sum compare success.." << std::endl;
  // Backward iteration test
  {
    InlineART<TestComparator>::Iterator iter(&list);
    Key start = 99999;
    iter.SeekForPrev(Encode(&start));

    int listSum = 0;
    while (iter.Valid()){
      listSum++;
      //Slice buf = GetLengthPrefixedSlice(iter.key());
      //Slice mvccKey = Slice(buf.data(), buf.size() - 8);
      //std::cout <<"index : " << listSum << ", key: " << mvccKey.data() <<std::endl;
      iter.Prev();
    }
    ASSERT_EQ(listSum, keys.size());
  }
  std::cout <<"sum BACK compare success.." <<std::endl;
}

TEST_F(InlineARTTest, InsertWithHint_Sequential) {
  const int N = 10000;
  const int S = 100;
  Arena arena;
  TestComparator cmp;
  TestInlineART list(cmp);
  void* hint = nullptr;
  for (int i = N; i < 2 * N; i++) {
    Key key = i;
    InsertWithHint(&list, key, &hint);
  }
  Validate(&list);
}

TEST_F(InlineARTTest, InsertWithHint_MultipleHints) {
  const int N = 10000;
  const int S = 100;
  Random rnd(534);
  Arena arena;
  TestComparator cmp;
  TestInlineART list(cmp);
  void* hints[S];
  Key last_key[S];
  for (int i = 0; i < S; i++) {
    hints[i] = nullptr;
    last_key[i] = 0;
  }
  for (int i = 0; i < N; i++) {
    Key s = rnd.Uniform(S);
    Key key = ((S + s) << 32) + (++last_key[s]);
    //std::cout << "index: " << i << ", key: " << key << std::endl;
    InsertWithHint(&list, key, &hints[s]);
  }
  Validate(&list);
  std::cout << "++++++++ InsertWithHint_MultipleHints  end .ok" <<std::endl;
}

TEST_F(InlineARTTest, InsertWithHint_MultipleHintsRandom) {
  const int N = 10000;
  const int S = 100;
  Random rnd(534);
  Arena arena;
  TestComparator cmp;
  TestInlineART list(cmp);
  void* hints[S];
  for (int i = 0; i < S; i++) {
    hints[i] = nullptr;
  }
  for (int i = 0; i <N; i++) {
    Key s = rnd.Uniform(S);
    Key key = ((s + S) << 32) + rnd.Next();
    InsertWithHint(&list, key, &hints[s]);
  }
  Validate(&list);
  std::cout << "InsertWithHint_MultipleHintsRandom end ok. " << std::endl;
}

TEST_F(InlineARTTest, InsertWithHint_CompatibleWithInsertWithoutHint) {
  const int N = 10000;
  const int S1 = 100;
  const int S2 = 100;
  Random rnd(534);
  Arena arena;
  TestComparator cmp;
  TestInlineART list(cmp);
  std::unordered_set<Key> used;
  Key with_hint[S1];
  Key without_hint[S2];
  void* hints[S1];
  for (int i = 0; i < S1; i++) {
    hints[i] = nullptr;
    while (true) {
      Key s = (rnd.Next() & 0x0FFFFFFF) | 0x20000000; // 保证所有key具有相同的长度，64位中前四位 决定长度区间
      if (used.insert(s).second) {
        with_hint[i] = s;
        break;
      }
    }
  }
  for (int i = 0; i < S2; i++) {
    while (true) {
      Key s = (rnd.Next() & 0x0FFFFFFF) | 0x20000000;
      if (used.insert(s).second) {
        without_hint[i] = s;
        break;
      }
    }
  }
  for (int i = 0; i < N; i++) {
    Key s = rnd.Uniform(S1 + S2);
    if (s < S1) {
      Key key = (with_hint[s] << 32) + rnd.Next();
      InsertWithHint(&list, key, &hints[s]);
    } else {
      Key key = (without_hint[s - S1] << 32) + rnd.Next();
      Insert(&list, key);
    }
  }
  Validate(&list);
  std::cout << "InsertWithHint_CompatibleWithInsertWithoutHint end. ok " <<std::endl;
}

#ifndef ROCKSDB_VALGRIND_RUN
// We want to make sure that with a single writer and multiple
// concurrent readers (with no synchronization other than when a
// reader's iterator is created), the reader always observes all the
// data that was present in the skip list when the iterator was
// constructor.  Because insertions are happening concurrently, we may
// also observe new values that were inserted since the iterator was
// constructed, but we should never miss any values that were present
// at iterator construction time.
//
// We generate multi-part keys:
//     <key,gen,hash>
// where:
//     key is in range [0..K-1]
//     gen is a generation number for key
//     hash is hash(key,gen)
//
// The insertion code picks a random key, sets gen to be 1 + the last
// generation number inserted for that key, and sets hash to Hash(key,gen).
//
// At the beginning of a read, we snapshot the last inserted
// generation number for each key.  We then iterate, including random
// calls to Next() and Seek().  For every key we encounter, we
// check that it is either expected given the initial snapshot or has
// been concurrently added since the iterator started.
class ConcurrentTest {
 public:
  static const uint32_t K = 8;

 private:
  static uint64_t key(Key key) { return (key >> 40); }
  static uint64_t gen(Key key) { return (key >> 8) & 0xffffffffu; }
  static uint64_t hash(Key key) { return key & 0xff; }

  static uint64_t HashNumbers(uint64_t k, uint64_t g) {
    uint64_t data[2] = {k, g};
    return Hash(reinterpret_cast<char*>(data), sizeof(data), 0);
  }

  static Key MakeKey(uint64_t k, uint64_t g) {
    assert(sizeof(Key) == sizeof(uint64_t));
    assert(k <= K);  // We sometimes pass K to seek to the end of the skiplist
    assert(g <= 0xffffffffu);
    if ( k == 0) 
      k = 1;
    return ((k << 40) | (g << 8) | (HashNumbers(k, g) & 0xff));
  }

  static bool IsValidKey(Key k) {
    return hash(k) == (HashNumbers(key(k), gen(k)) & 0xff);
  }

  static Key RandomTarget(Random* rnd) {
    switch (rnd->Next() % 10) {
      case 0:
        // Seek to beginning
        return MakeKey(0, 0);
      case 1:
        // Seek to end
        return MakeKey(K, 0);
      default:
        // Seek to middle
        return MakeKey(rnd->Next() % K, 0);
    }
  }

  // Per-key generation
  struct State {
    std::atomic<int> generation[K];
    void Set(int k, int v) {
      generation[k].store(v, std::memory_order_release);
    }
    int Get(int k) { return generation[k].load(std::memory_order_acquire); }

    State() {
      for (unsigned int k = 0; k < K; k++) {
        Set(k, 0);
      }
    }
  };

  // Current state of the test
  State current_;

  ConcurrentArena arena_;

  // InlineART is not protected by mu_.  We just use a single writer
  // thread to modify it.
  InlineART<TestComparator> list_;

 public:
  ConcurrentTest() : list_(TestComparator()) {}

  // REQUIRES: No concurrent calls to WriteStep or ConcurrentWriteStep
  void WriteStep(Random* rnd) {
    uint32_t k = rnd->Next() % K;
    k = (k == 0 ? 1 : k);
    const int g = current_.Get(k) + 1;
    const Key new_key = MakeKey(k, g);
    char* buf = genKey(list_, new_key);
    list_.Insert(buf);
    current_.Set(k, g);
  }

  // REQUIRES: No concurrent calls for the same k
  void ConcurrentWriteStep(uint32_t k) {
    if (k == 0) k = 1;
    const int g = current_.Get(k) + 1;
    const Key new_key = MakeKey(k, g);
    char* buf = genKey(list_, new_key);
    list_.Insert(buf);
    ASSERT_LE(g, current_.Get(k) + 1);
    current_.Set(k, current_.Get(k) + 1);
  }

  void ReadStep(Random* rnd) {
    // Remember the initial committed state of the skiplist.
    State initial_state;
    for (unsigned int k = 0; k < K; k++) {
      initial_state.Set(k, current_.Get(k));
    }

    Key pos = RandomTarget(rnd);
    InlineART<TestComparator>::Iterator iter(&list_);
    iter.Seek(Encode(&pos));
    while (true) {
      Key current;
      if (!iter.Valid()) {
        current = MakeKey(K, 0);
      } else {
        current = Decode(iter.key());
        ASSERT_TRUE(IsValidKey(current)) << current;
      }
      if (pos > current){
        std::cout <<"ERROR , cannot be here . pos: "<< pos << ", current: " << current << std::endl;
        exit(0);
      }
      ASSERT_LE(pos, current) << "should not go backwards";
      

      // Verify that everything in [pos,current) was not present in
      // initial_state.
      while (pos < current) {
        ASSERT_LT(key(pos), K) << pos;
        // Advance to next key in the valid key space
        if (key(pos) < key(current)) {
          pos = MakeKey(key(pos) + 1, 0);
        } else {
          pos = MakeKey(key(pos), gen(pos) + 1);
        }
      }

      if (!iter.Valid()) {
        break;
      }

      if (rnd->Next() % 2) {
        iter.Next();
      } else {
        Key new_target = RandomTarget(rnd);
        if (new_target > pos) {
          pos = new_target;
          iter.Seek(Encode(&new_target));
        }
      }
    }
  }
};
const uint32_t ConcurrentTest::K;

// Simple test that does single-threaded testing of the ConcurrentTest
// scaffolding.
TEST_F(InlineARTTest, ConcurrentReadWithoutThreads) {
  ConcurrentTest test;
  Random rnd(test::RandomSeed());
  for (int i = 0; i < 10000; i++) {
    test.ReadStep(&rnd);
    test.WriteStep(&rnd);
  }
  std::cout <<"ConcurrentReadWithoutThreads end. "<<std::endl;
}

TEST_F(InlineARTTest, ConcurrentInsertWithoutThreads) {
  ConcurrentTest test;
  Random rnd(test::RandomSeed());
  for (int i = 0; i < 10000; i++) {
    test.ReadStep(&rnd);
    uint32_t base = rnd.Next();
    for (int j = 0; j < 4; ++j) {
      test.ConcurrentWriteStep((base + j) % ConcurrentTest::K);
    }
  }
  std::cout <<"ConcurrentInsertWithoutThreads end. "<<std::endl;
}

class TestState {
 public:
  ConcurrentTest t_;
  int seed_;
  std::atomic<bool> quit_flag_;
  std::atomic<uint32_t> next_writer_;

  enum ReaderState { STARTING, RUNNING, DONE };

  explicit TestState(int s)
      : seed_(s),
        quit_flag_(false),
        state_(STARTING),
        pending_writers_(0),
        state_cv_(&mu_) {}

  void Wait(ReaderState s) {
    mu_.Lock();
    while (state_ != s) {
      state_cv_.Wait();
    }
    mu_.Unlock();
  }

  void Change(ReaderState s) {
    mu_.Lock();
    state_ = s;
    state_cv_.Signal();
    mu_.Unlock();
  }

  void AdjustPendingWriters(int delta) {
    mu_.Lock();
    pending_writers_ += delta;
    if (pending_writers_ == 0) {
      state_cv_.Signal();
    }
    mu_.Unlock();
  }

  void WaitForPendingWriters() {
    mu_.Lock();
    while (pending_writers_ != 0) {
      state_cv_.Wait();
    }
    mu_.Unlock();
  }

 private:
  port::Mutex mu_;
  ReaderState state_;
  int pending_writers_;
  port::CondVar state_cv_;
};

static void ConcurrentReader(void* arg) {
  TestState* state = reinterpret_cast<TestState*>(arg);
  Random rnd(state->seed_);
  int64_t reads = 0;
  state->Change(TestState::RUNNING);
  while (!state->quit_flag_.load(std::memory_order_acquire)) {
    state->t_.ReadStep(&rnd);
    ++reads;
  }
  state->Change(TestState::DONE);
}

static void ConcurrentWriter(void* arg) {
  TestState* state = reinterpret_cast<TestState*>(arg);
  uint32_t k = state->next_writer_++ % ConcurrentTest::K;
  //std::cout << "ConcurrentWriter k: " << k << std::endl;
  state->t_.ConcurrentWriteStep(k);
  state->AdjustPendingWriters(-1);
}

static void RunConcurrentRead(int run) {
  const int seed = test::RandomSeed() + (run * 100);
  Random rnd(seed);
  const int N = 1000;
  const int kSize = 1000;
  for (int i = 0; i < N; i++) {
    if ((i % 100) == 0) {
      fprintf(stderr, "Run %d of %d\n", i, N);
    }
    TestState state(seed + 1);
    Env::Default()->SetBackgroundThreads(1);
    Env::Default()->Schedule(ConcurrentReader, &state);
    state.Wait(TestState::RUNNING);
    for (int k = 0; k < kSize; ++k) {
      state.t_.WriteStep(&rnd);
    }
    state.quit_flag_.store(true, std::memory_order_release);
    state.Wait(TestState::DONE);
  }
}

static void RunConcurrentInsert(int run, int write_parallelism = 4) {
  Env::Default()->SetBackgroundThreads(1 + write_parallelism,
                                       Env::Priority::LOW);
  const int seed = test::RandomSeed() + (run * 100);
  Random rnd(seed);
  const int N = 1000;
  const int kSize = 1000;
  for (int i = 0; i < N; i++) {
    if ((i % 100) == 0) {
      fprintf(stderr, "RunConcurrentInsert  Run %d of %d\n", i, N);
    }
    TestState state(seed + 1);
    //Env::Default()->Schedule(ConcurrentReader, &state);
    //state.Wait(TestState::RUNNING);
    for (int k = 0; k < kSize; k += write_parallelism) {
      state.next_writer_ = rnd.Next();
      state.AdjustPendingWriters(write_parallelism);
      for (int p = 0; p < write_parallelism; ++p) {
        Env::Default()->Schedule(ConcurrentWriter, &state);
      }
      state.WaitForPendingWriters();
    }
    state.quit_flag_.store(true, std::memory_order_release);
    //state.Wait(TestState::DONE);
  }
}

TEST_F(InlineARTTest, ConcurrentRead1) { RunConcurrentRead(1); }
TEST_F(InlineARTTest, ConcurrentRead2) { RunConcurrentRead(2); }
TEST_F(InlineARTTest, ConcurrentRead3) { RunConcurrentRead(3); }
TEST_F(InlineARTTest, ConcurrentRead4) { RunConcurrentRead(4); }
TEST_F(InlineARTTest, ConcurrentRead5) { RunConcurrentRead(5); }
TEST_F(InlineARTTest, ConcurrentInsert1) { RunConcurrentInsert(1); }
TEST_F(InlineARTTest, ConcurrentInsert2) { RunConcurrentInsert(2); }
TEST_F(InlineARTTest, ConcurrentInsert3) { RunConcurrentInsert(3); }

#endif  // ROCKSDB_VALGRIND_RUN
}  // namespace rocksdb

int main(int argc, char** argv) {
  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}
